#include<iostream>
#include<fstream>
#include<sstream>
#include<cstdlib>
#include<list>
#include<string.h>
#include<map>

#include<gsl/gsl_blas.h>
#include<gsl/gsl_linalg.h>
#include<gsl/gsl_matrix.h>
#include<gsl/gsl_multifit.h>
#include<gsl/gsl_permutation.h>
#include<gsl/gsl_randist.h>
#include<gsl/gsl_cdf.h>
#include<gsl/gsl_statistics_double.h>
#include"../common/main.hpp"
#include"../common/analyzer.hpp"
#include"../common/io.hpp"
#include"univariate.hpp"

using namespace std;

typedef short unsigned int suint;

void Univariate::init(const ptree & pt){
  settings = new univariate_settings_t;
  settings->phenofile = pt.get<string>("phenofile");
  settings->maskfile = pt.get<string>("maskfile");
  settings->annotationfile = pt.get<string>("annotationfile");
  settings->type = pt.get<string>("type");
  settings->regression = pt.get<string>("regression");
  settings->geno_format = pt.get<char>("geno_format");
  mask = NULL;
}

Univariate::Univariate(){}

Univariate::~Univariate(){
  if (mask!= NULL){
     delete[] mask;
  }
  cerr<<"Univariate object deleted\n";
}

void Univariate::getDouble(char t_i,const string & str, double * output){
  uint j=0;
  switch (t_i){
    case 'S':
      for(uint i=0;i<veclen;++i){
        if (mask[i]){
          switch(str[i]){
            //case '9':
            case '0':
              output[j] = 0;
              break;
            case '1':
              output[j] = 1;
              break;
            case '2':
              output[j] = 2;
              break;
          }
          //cerr<<output[j];
          ++j;
        }
      }
      break;
    case 'T':
      for(uint i=0;i<veclen;++i){
        if (mask[i]){
          switch(str[i]){
            //case '9':
            case '0':
              output[j] = IO::MISSING_GENO;
              break;
            //case '0':
            case '1':
              output[j] = 0;
              break;
            //case '1':
            case '2':
              output[j] = 1;
              break;
            //case '2':
            case '3':
              output[j] = 2;
              break;
          }
          ++j;
        }
      }
      break;
    case 'I':
      for(uint i=0;i<veclen;++i){
        if (mask[i]){
          output[j++] = ((int)str[i]-65)/10.;
        }
      }
      break;

  }
}

void Univariate::loadmask(){
    string line;
    ifstream maskFile(MASKFILE);
    bool temp[10000];
    veclen = 0;
    if (maskFile){
      cerr<<"Found a geno mask file.\n";
      if (maskFile.is_open()) {
        veclen = 0;
        while(getline (maskFile,line)){
          istringstream iss(line);
          bool m;
          iss>>m;
          temp[veclen++]=m;
        }
        maskFile.close();
      }
    }else{
      cerr<<"No mask file found.\n";
      veclen = persons;
      for(unsigned int i=0;i<veclen;++i){
        ostringstream oss;
        oss<<i;
        temp[i]=true;
      }
    }
    mask = new bool[veclen];
    for(uint i=0;i<veclen;++i){
      mask[i] = temp[i];
    }
    cerr<<"Length of mask: "<<veclen<<endl;
}

double  Univariate::HWEpvalue(uint * cellcounts,double maf){
    if (maf==1.||maf==0.) return 0;
    double other = 1.-maf;
    uint total=0;
    for(uint i=0;i<3;++i){
        total+=cellcounts[i];
    }
    if (!total) return 0;

    double exp[3];
    exp[0] = total*1.*other*other;
    exp[1] = total*2.*maf*other;
    exp[2] = total*1.*maf*maf;
    double chi=0;
//    cerr<<"MAF: "<<maf<<", total: "<<total<<endl;
    for(uint i=0;i<3;++i){
//        cerr<<"Obs: "<<cellcounts[i]<<endl;
//        cerr<<"Exp: "<<exp[i]<<endl;
        chi+=pow(cellcounts[i]-exp[i],2)/exp[i];
    }
//cerr<<"chi is "<<chi<<endl;
    double pval = 1.-gsl_cdf_chisq_P(chi,1);

    return pval;
}

void Univariate::printtyped(double * geno,double & callrate,double & maf, double & hwe, double & exp_het, double & obs_het){
  uint callcounts1=0,totalallelecount1=0,allelecount1=0,heterocounts1=0;
  uint cellCounts[3];
  for(uint j=0;j<3;++j) cellCounts[j] = 0;
  for(uint j=0;j<persons;++j){
    bool call1 = (geno[j]!=IO::MISSING_GENO)?true:false;
    int call1geno = static_cast<int>(geno[j]);
    callcounts1+=call1;
    // BEGIN HWE CALCULATIONS
    if (geno[j]!=IO::MISSING_GENO) {
      ++cellCounts[call1geno];
      totalallelecount1+=2;
      allelecount1+=call1geno;
    }
    heterocounts1+=(call1 && call1geno==1)?1:0;
    // END HWE CALCULATIONS
  }
  callrate = (callcounts1*1./persons);
  maf=allelecount1*1./totalallelecount1;
  hwe = HWEpvalue(cellCounts,maf);
  exp_het = (2.*maf*(1-maf));
  obs_het = (1.*heterocounts1/callcounts1);
//  cout<<"\t"<<(callcounts1*1./persons)
//        <<"\t"<<maf1
//        <<"\t"<<hwe1
//        <<"\t"<<(1.*heterocounts1/callcounts1)
//        <<"\t"<<(2.*maf1*(1-maf1));
}

void Univariate::printimputed(double * geno, double & maf1){
  double genoFreq = 0.;
  for(uint j=0;j<persons;++j){
    genoFreq+=geno[j];
  }
  maf1=(genoFreq/(persons*1.0)/2.0);

}

void Univariate::loadAnnotation(){
  ifstream ifs(ANNOTATION_FILE);
  if (ifs.is_open()){
    cerr<<"Annotation file found. Loading...\n";
    printAnnot = true;
    string line;
    getline(ifs,line);
    while(getline(ifs,line)){
      istringstream iss(line);
      string rs;
      //short int chr;
      //int position;
      //char other,ref,type;
      char type;
      //float rsq; 
      snpinfo s;
      iss>>rs>>s.chr>>s.pos>>s.other>>s.ref>>type>>s.rsq;
      //chr_map[rs]=chr;
      //pos_map[rs]=position;
      //other_map[rs]=other;
      //ref_map[rs]=ref;
      //rsq_map[rs]=rsq;
      snpinfo_map[rs] = s;
    }
    ifs.close();
    cerr<<"Annotation loaded.\n";
  }else{
    cerr<<"Cannot find "<<ANNOTATION_FILE<<endl;
    printAnnot = false;
  }
}

void Univariate::loadPheno(double * phenovec,double * designmat, bool * defined,uint & rank ){
  //phenovec = new double[persons];
  //defined = new bool[persons];
  ifstream ifs_pheno(phenofile);
  if (!ifs_pheno.is_open()) throw "Can't open phenotypes.txt";
  // extra 3 are localanc,genotype,and interaction
  //rank = marginal?covariates+1:covariates+3;
  //cerr<<"Allocating a design matrix of "<<rank<<" cols\n";
  //designmat = new double[persons*rank];
  string line;
  for(uint i=0;i<persons;++i){
     getline(ifs_pheno,line);
//cerr<<line<<endl;
     istringstream iss(line);
     designmat[i*rank] = 1.;
     defined[i] = true;
     for(uint j=0;j<covariates;++j){
       double val;
       iss>>val;
        //cerr<<"cov:"<<j<<" person:"<<i<<":"<<val<<endl;
       if (fabs(val-IO::AFF_UNDEF)<1) {
        //cerr<<"person "<<i<<" missing covariate\n";
        defined[i] = false;
       }
       //else{
         if (j){
           designmat[i*rank+j] = val;
            //cerr<<i<<" col :"<<j<<" "<<designmat[i*rank+j]<<endl;
         }else{
           phenovec[i] = val;
            //cerr<<i<<":"<<phenovec[i]<<endl;
         }
       //}
     }
     //cerr<<phenovec[i];
  }
  ifs_pheno.close();
  //bool * mask;
}


void Univariate::run(){
  phenofile = settings->phenofile.data();
  MASKFILE = settings->maskfile.data();
  ANNOTATION_FILE = settings->annotationfile.data();
  logistic = !settings->regression.compare("logistic")?true:false;
  marginal = !settings->type.compare("marginal")?true:false;
  typed_imputed = settings->geno_format;
  //infer dimensions
  persons = 1;
  covariates = 0;
  uint rank = 0;
  ifstream ifs(phenofile);
  if (!ifs.is_open()){
    throw "Cannot open phenotype file";
  }
  string line,token;
  getline(ifs,line);
  istringstream iss(line);
  while(iss>>token) ++covariates;
  cerr<<"covariates: "<<covariates<<endl;
  while(getline(ifs,line)) ++persons;
  ifs.close();
  cerr<<"person: "<<persons<<endl;
  double phenovec[persons];
  bool defined[persons];
  // extra 3 are localanc,genotype,and interaction
  rank = marginal?covariates+1:covariates+3;
  cerr<<"Allocating a design matrix of "<<rank<<" cols\n";
  double designmat [persons*rank];

  
  bool cachedmodel = false;
  double L0,L1;
  bool imputed;
  switch(typed_imputed){
  case 'T':
    imputed = false;
    break;
  case 'I':
    imputed = true;
    break;
  default:
    cerr<<"Typed imputed code not recognized: "<<typed_imputed<<endl;
    exit(0);
    break;
  }
 
  loadAnnotation();

  //uint veclen;
  loadmask();
  cerr<<"Mask loaded\n";
  loadPheno(phenovec, designmat, defined,rank);
  cerr<<"Phenotypes loaded\n";
  //cout<<"SNP\tCHI_LOCALANC\tCHI_GENO\tCHI_INTER\n";
  long unsigned int row=1;

  cerr<<"Iteration:\n";
  if (printAnnot){
    cout<<"SNPID\tchr\tposition\tcoded_all\tnoncoded_all\tstrand_genome\tbeta\tSE\tpval\tAF_coded_all\tHWE_pval\tcallrate\tn_total\timputed\tused_for_imp\toevar_imp\tMAC\n";
  }else{
    if (imputed){
      ofs_snpstats.open("report.imputed.snp");
      ofs_snpstats<<"SNP\tMAF\n";
    }else{
      ofs_snpstats.open("report.snp");
      ofs_snpstats<<"SNP\tCallrate\tMAF\tHWE\tHet_Obs\tHet_Exp"<<endl;
    }
  }

  while(getline(cin,line)){
    double betas[rank];
    memset(betas,0,sizeof(betas));
    double var[rank*rank];
    if ((row % 100) == 0) cerr<<" "<<row;

    istringstream iss(line);
    string rs,local,geno;
    if (marginal){
      iss>>rs>>geno;
    }else{
      iss>>rs>>typed_imputed>>local>>geno;
    }
    snpinfo s;
    if (printAnnot){
      s = fetch(snpinfo_map,rs);
    }
    double d_local[persons];
    getDouble('S',local,d_local);
    double d_geno[persons];
    getDouble(typed_imputed,geno,d_geno);
    double callrate=1,maf=1,hwe=1,exp_het=0,obs_het=0;
    switch(typed_imputed){
      case 'T':
        printtyped(d_geno,callrate,maf,hwe,exp_het,obs_het);
        break;
      case 'I':
        printimputed(d_geno,maf);
        break;
      default:
        cerr<<"Typed imputed code not recognized: "<<typed_imputed<<endl;
        exit(0);
        break;
    }
   
    uint samplesize=0;
    int affvec_filtered[persons];
    float designmat_filtered[persons * rank];

    double phenovec_filtered[persons];
    double designmat_filtered_double[persons * rank];
    // Load the design matrix
    //ofstream ofs("debug");
    for(uint person=0;person<persons;++person){
      if (!defined[person]||d_geno[person] == IO::MISSING_GENO ){
      //cerr<<"person "<<person<<" missing geno or cov\n";
      }else{
        phenovec_filtered[samplesize] = phenovec[person];
        affvec_filtered[samplesize] =  static_cast<int>(phenovec_filtered[samplesize]);
        for(uint j=0;j<covariates;++j){
          designmat_filtered[samplesize*rank+j] = designmat_filtered_double[samplesize*rank+j] = designmat[person*rank+j];
        }
        if (marginal){
          designmat_filtered[samplesize*rank+covariates] = designmat_filtered_double[samplesize*rank+covariates] = d_geno[person];
        }else{
          designmat_filtered[samplesize*rank+covariates] = designmat_filtered_double[samplesize*rank+covariates] = d_local[person];
          designmat_filtered[samplesize*rank+covariates+1] = designmat_filtered_double[samplesize*rank+covariates+1] = d_geno[person];
          designmat_filtered[samplesize*rank+covariates+2] = designmat_filtered_double[samplesize*rank+covariates+2] = d_local[person] * d_geno[person];
        }
        ++samplesize;
      }
    }
    //ofs.close();
    //exit(0);
    double mac = (maf<.5)?samplesize*maf*1.:samplesize*(1.-maf)*1.;
    // run linear regression
    double dChiSq=0.,pvalue=0.;
    bool fitted = false;
    if (logistic){
      if (!imputed || !cachedmodel || !marginal){
        fitModel(L0,affvec_filtered,designmat_filtered,betas,var,
        samplesize,rank,rank-1);
        cachedmodel = true;
      }
      //fitted = fitModelCCD(L1,affvec_filtered,designmat_filtered, betas,var,samplesize,rank,rank);
      //cerr<<"CCD test:";
      //for(uint j=0;j<rank;++j){
       // cerr<<" "<<betas[j];
        //betas[j] = 0.;
      //}
      //cerr<<endl;
      fitted = fitModel(L1,affvec_filtered,designmat_filtered, betas,var,samplesize,rank,rank);
      //cerr<<"logistic test:";
      //for(uint j=0;j<rank;++j){
        //cerr<<" "<<betas[j];
        //betas[j] = 0.;
      //}
      //cerr<<endl;
      if (fitted){
        dChiSq = 2.*(L1-L0);
        pvalue = (1.-gsl_cdf_chisq_P(dChiSq,1));
      }
    }else{
      double  resid[samplesize];
      fitted =  leastsquares(designmat_filtered_double, phenovec_filtered,betas, var,resid, rank, samplesize);
      if (fitted){
        double wald2 = pow(betas[rank-1],2) / var[rank*rank-1];
        double mean = gsl_stats_mean(phenovec_filtered,1,samplesize);
        double tss=0.;
        double rss=0.;
        for(uint j=0;j<samplesize;++j){
          rss+=pow(resid[j],2);
          tss+=pow(phenovec_filtered[j]-mean,2);
        }
        dChiSq=1.-rss/tss;
        pvalue =  (1.-gsl_cdf_chisq_P(wald2,1));
      }
    }
    if (fitted){
      if (printAnnot){
        cout<<rs
        <<"\t"<<s.chr
        <<"\t"<<s.pos
        <<"\t"<<s.other
        <<"\t"<<s.ref
        <<"\t+"
        <<"\t"<<betas[rank-1]
        <<"\t"<<sqrt(var[(rank*rank)-1])
        <<"\t"<<pvalue
        <<"\t"<<maf
        <<"\t"<<hwe
        <<"\t"<<callrate
        <<"\t"<<samplesize
        <<"\t"<<imputed
        <<"\t"<<!imputed
        <<"\t"<<s.rsq
        <<"\t"<<mac
        <<endl;
      }else{
        cout<<rs
        <<"\t"<<betas[rank-1]
        <<"\t"<<sqrt(var[(rank*rank)-1])
        <<"\t"<<dChiSq
        <<"\t"<<pvalue
        <<endl;
        if (imputed){
          ofs_snpstats<<rs<<"\t"<<maf<<endl;
        }else{
          ofs_snpstats<<rs
          <<"\t"<<callrate
          <<"\t"<<maf
          <<"\t"<<hwe
          <<"\t"<<obs_het
          <<"\t"<<exp_het
         <<endl;
        }
      }
    }
    ++row;
  }  
  cerr<<"\nComplete:\n";
  if (!printAnnot){
    ofs_snpstats.close();
  }
}
